import cv2
import numpy as np
import matplotlib.pyplot as plt

# 生成运动模糊核
def make_PSE(kernel_size=15, angle=60):
    PSE = np.diag(np.ones(kernel_size))  # 创建初始模糊核，方向为 -45°
    angle += 45  # 调整角度，抵消初始 -45° 的影响
    # 生成旋转矩阵并应用于模糊核
    M = cv2.getRotationMatrix2D((kernel_size / 2, kernel_size / 2), angle, 1)
    PSE = cv2.warpAffine(PSE, M, (kernel_size, kernel_size))
    PSE = PSE / PSE.sum()  # 权重归一化，确保模糊核的总和为1
    return PSE

# 在空域对图像进行运动模糊
def motion_blur(image, PSE):
    blurred = cv2.filter2D(image, -1, PSE)  # 使用 2D 卷积进行模糊处理
    # 对模糊后的图像进行归一化处理，确保像素值在 [0, 255] 范围内
    cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
    return np.array(blurred, dtype=np.uint8)

# 添加标准高斯噪声
def make_noisy(input_img, c=0.1):
    noise = c * input_img.std() * np.random.normal(loc=0.0, scale=1.0, size=input_img.shape)
    noisy_img = input_img + noise  # 将噪声添加到原图像上
    noisy_img = np.clip(noisy_img, 0, 255)  # 将像素值限制在 [0, 255] 范围内
    return noisy_img.astype(np.uint8)

# 逆滤波函数
def inverse_filtering(input_img, PSF, eps=1e-5):
    img_h, img_w = input_img.shape[:2]
    # 将 PSF 填充到与输入图像相同大小
    PSF_padded = np.zeros_like(input_img, dtype=np.float32)
    kh, kw = PSF.shape
    PSF_padded[:kh, :kw] = PSF

    input_fft = np.fft.fft2(input_img)  # 对输入图像进行 FFT 变换
    PSF_fft = np.fft.fft2(PSF_padded) + eps  # 对 PSF 进行 FFT，并添加小值以防分母为零

    # 执行逆滤波：在频域上除以 PSF 的频谱
    result_fft = input_fft / PSF_fft
    result = np.abs(np.fft.ifft2(result_fft))  # 逆 FFT 还原到空域
    return np.clip(result, 0, 255).astype(np.uint8)

# 维纳滤波函数
def wiener_filter(input_img, PSF, K=0.01):
    img_h, img_w = input_img.shape[:2]
    # 将 PSF 填充到与输入图像相同大小
    PSF_padded = np.zeros_like(input_img, dtype=np.float32)
    kh, kw = PSF.shape
    PSF_padded[:kh, :kw] = PSF

    input_fft = np.fft.fft2(input_img)  # 对输入图像进行 FFT 变换
    PSF_fft = np.fft.fft2(PSF_padded)  # 对 PSF 进行 FFT

    # 维纳滤波公式：利用 PSF 的共轭、功率谱及噪声功率比
    PSF_fft_conj = np.conj(PSF_fft)
    wiener_filter = PSF_fft_conj / (np.abs(PSF_fft)**2 + K)  # 计算维纳滤波器

    result_fft = input_fft * wiener_filter  # 频域滤波
    result = np.abs(np.fft.ifft2(result_fft))  # 逆 FFT 还原到空域
    return np.clip(result, 0, 255).astype(np.uint8)

# 主程序
image_path = "img/kenny.jpg"
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

# 生成模糊核并应用运动模糊
kernel_size = 15
angle = 60
PSE = make_PSE(kernel_size, angle)
blurred_image = motion_blur(image, PSE)

# 添加噪声
noisy_blurred_image = make_noisy(blurred_image, c=0.1)

# 使用逆滤波复原
restored_inverse = inverse_filtering(noisy_blurred_image, PSE)

# 使用维纳滤波复原
restored_wiener = wiener_filter(noisy_blurred_image, PSE, K=0.01)

# 显示结果
plt.figure(figsize=(10, 8))

plt.subplot(2, 2, 1)
plt.title("Original Image")
plt.imshow(image, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 2)
plt.title("Motion Blurred Image with Noise")
plt.imshow(noisy_blurred_image, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 3)
plt.title("Inverse Filter Restored")
plt.imshow(restored_inverse, cmap='gray')
plt.axis('off')

plt.subplot(2, 2, 4)
plt.title("Wiener Filter Restored")
plt.imshow(restored_wiener, cmap='gray')
plt.axis('off')

plt.tight_layout()
plt.show()
